'''
Author: Yicheng Chen (yicheng-chen@outlook.com)
LastEditTime: 2023-05-24 10:01:05
'''
import os
import sys
current_path = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, current_path)
from visualizer import Visualizer
from visualization_msgs.msg import MarkerArray,Marker
import rospy
import numpy as np
from mavros_msgs.msg import State, PositionTarget
from mavros_msgs.srv import SetMode, SetModeRequest
from geometry_msgs.msg import PoseStamped
from traj_planner import MinJerkPlanner
from pyquaternion import Quaternion
import time
from nav_msgs.msg import Odometry, Path
import actionlib
from traj_generator.msg import *
import astar


class MissionConfig():
    def __init__(self):
        self.des_pos_z = rospy.get_param("~des_pos_z", 2.0)
        self.target_reach_threshold = rospy.get_param("~target_reach_threshold", 0.2)
        self.cmd_hz = rospy.get_param("~cmd_hz", 300)


class DroneState():
    def __init__(self):
        self.global_pos = np.zeros(3)
        self.global_vel = np.zeros(3)
        self.local_vel = np.zeros(3)
        self.attitude = Quaternion()


class TrajPlanner():
    def __init__(self, node_name="traj_planner"):
        # Node
        rospy.init_node(node_name, anonymous=False)

        # Members
        self.des_path = Path()
        self.visualizer = Visualizer()
        self.drone_state = DroneState()
        mission_config = MissionConfig()
        self.planner = MinJerkPlanner()
        self.state_cmd = PositionTarget()
        self.state_cmd.coordinate_frame = 1

        # Parameters
        self.des_pos_z = mission_config.des_pos_z
        self.target_reach_threshold = mission_config.target_reach_threshold
        self.cmd_hz = mission_config.cmd_hz
        self.planning_hz=0.2
        self.T = 30
        self.start_px = 0
        self.start_py = 0
        self.end_px=10
        self.end_py=90
        self.grid_resolution = 1.0
        self.ox_centre=7
        self.oy_centre=35
        self.obstacle_r=2
        self.ox, self.oy = astar.get_env(self.ox_centre,self.oy_centre,self.obstacle_r)
        # self.svx=0
        # self.svy=0
        self.radius = 2 
        self.pathx=[]
        self.pathy=[]

        # Flags and counters
        self.mission_executing = False
        self.near_global_target = False
        self.reached_target = False
        self.odom_received = False
        self.des_state_index = 0
        self.future_index = 99999
        self.plan_times=0
        self.over_planning_flag =0
        

        # Server
        self.plan_server = actionlib.SimpleActionServer('plan', PlanAction, self.execute_mission, False)
        self.plan_server.start()

        # Services
        rospy.wait_for_service("/mavros/set_mode")
        self.set_mode_client = rospy.ServiceProxy("mavros/set_mode", SetMode)

        # Subscribers
        self.flight_state_sub = rospy.Subscriber('/mavros/state', State, self.flight_state_cb)
        self.odom_sub = rospy.Subscriber('/mavros/local_position/odom', Odometry, self.odom_cb)
        self.goal_sub = rospy.Subscriber('goal_pose', PoseStamped, self.goal_pos_callback)
        self.obstacle_sub = rospy.Subscriber('obstacle_pose', PoseStamped, self.obstacle_pos_callback)

        # Publishers
        self.local_pos_cmd_pub = rospy.Publisher("/mavros/setpoint_raw/local", PositionTarget, queue_size=10)
        self.des_wpts_pub = rospy.Publisher('/des_wpts', MarkerArray, queue_size=10)
        self.des_path_pub = rospy.Publisher('/des_path', MarkerArray, queue_size=10)
        self.obstacle_region_pub = rospy.Publisher('/obstacle_region', Marker, queue_size=10)

        rospy.loginfo("Global planner initialized")

    def flight_state_cb(self, data):
        self.flight_state = data

    def odom_cb(self, data):
        '''
        1. store the drone's global status
        2. publish dynamic tf transform from map frame to camera frame
        (Currently, regard camera frame as drone body frame)
        '''
        self.odom_received = True
        self.odom = data
        local_pos = np.array([data.pose.pose.position.x,
                              data.pose.pose.position.y,
                              data.pose.pose.position.z])
        global_pos = local_pos
        local_vel = np.array([data.twist.twist.linear.x,
                              data.twist.twist.linear.y,
                              data.twist.twist.linear.z])
        quat = Quaternion(data.pose.pose.orientation.w,
                          data.pose.pose.orientation.x,
                          data.pose.pose.orientation.y,
                          data.pose.pose.orientation.z)  # from local to global
        global_vel = quat.rotate(local_vel)
        self.drone_state.global_pos = global_pos
        self.drone_state.global_vel = global_vel
        self.drone_state.local_vel = local_vel
        self.drone_state.attitude = quat

        if self.mission_executing and np.linalg.norm(global_pos[:2] - self.global_target) < self.target_reach_threshold:
            rospy.loginfo("Global target reached!\n")
            self.end_mission()

    def init_mission(self):
        self.mission_executing = True
        self.reached_target = False
        self.des_state_index = 0

    def end_mission(self):
        self.tracking_cmd_timer.shutdown()
        rospy.loginfo("over planning!")
        self.planning_timer.shutdown()
        self.mission_executing = False
        self.reached_target = True
        # self.des_state_index = 0

    def execute_mission(self, goal):
        # convert goal.int_wpts to numpy array
        self.target_wpts = np.array([goal.wpts]).reshape(-1, 3)
        self.global_target = self.target_wpts[-1, :2]
        self.init_mission()
        
        self.planning_timer = rospy.Timer(rospy.Duration(1/self.planning_hz), self.planning_timer_cb)
        self.report_planning_result()
        # self.gen_traj()

        # while not self.odom_received:
        #     time.sleep(0.01)

        # self.start_tracking()
        # self.visualize_des_path()
        # self.report_planning_result()

    def goal_pos_callback(self,point):
        self.end_px=point.pose.position.x
        self.end_py=point.pose.position.y
        rospy.loginfo("目标x的数据:%s",point.pose.position.x)
        
    def obstacle_pos_callback(self,point):
        self.ox_centre=point.pose.position.x
        self.oy_centre=point.pose.position.y
    
    def planning_timer_cb(self,event):

        if self.T>0.001:
            rospy.loginfo("T值:%s",self.T)
            self.gen_traj()
            self.T=self.T-1/self.planning_hz
            while not self.odom_received:
                time.sleep(0.01)
            self.start_tracking()
            self.visualize_obstacle()
            self.visualize_des_path()
            
        else:
            rospy.loginfo("over planning!")
            self.over_planning_flag=1
            return
        

    def gen_traj(self):
        self.plan_times+=1
        self.des_pos_z = self.target_wpts[0, 2]
        # 前端A*规划
        start_time = time.time()
        if self.plan_times ==1:
            self.start_px=2
            self.start_py=2
        else:
            self.start_px=self.drone_state.global_pos[0]
            self.start_py=self.drone_state.global_pos[1]
        print(self.start_px)
        self.end_px=self.target_wpts[-1,0]
        self.end_py=self.target_wpts[-1,1]
        print(self.end_px)
        self.ox=[]
        self.oy=[]
        self.ox,self.oy=astar.get_env(self.ox_centre,self.oy_centre,self.obstacle_r)
        self.pathx, self.pathy  = astar.astar_planning(self.start_px, self.start_py, self.end_px, self.end_py, self.ox, self.oy, self.grid_resolution, self.radius)
        self.pathx.pop()
        self.pathy.pop()
        # 后段Minimumjerk轨迹优化
        target_wpts_num=len(self.pathx)
        target_wpts_z=np.ones(target_wpts_num)*self.des_pos_z 
        target_wpts_2d=np.vstack((self.pathx,self.pathy))
        # target_wpts_2d=np.transpose(target_wpts_2d)
        self.target_wpts=[]
        self.target_wpts= np.vstack((target_wpts_2d,target_wpts_z))
        self.target_wpts=np.transpose(self.target_wpts)
        head_state = np.array([self.drone_state.global_pos[:2],
                            self.drone_state.global_vel[:2],
                            np.zeros(2)])
        tail_state = np.array([self.target_wpts[-1, :2],
                            np.zeros(2),
                            np.zeros(2)])
        if self.plan_times ==1:
            int_wpts = self.target_wpts[:-1, :2]
        else:
            int_wpts = self.target_wpts[1:-1, :2]
        # ts = self.T * np.ones(len(int_wpts)+1)  # you can change this to other time allocation method
        ts=self.arangeT()

        self.planner=MinJerkPlanner()
        self.planner.plan(3, head_state, tail_state, int_wpts, ts)
        end_time = time.time()

        self.des_state_array, self.traj_time, self.cmd_hz = self.planner.get_full_state_cmd()

        rospy.loginfo(f"Trajectory generated. Time cost: {end_time - start_time}s")

           

    def arangeT(self):
        x0=self.pathx[:-1]
        x1=self.pathx[1:]
        dis_x=np.array(x1)-np.array(x0)
        y0=self.pathy[:-1]
        y1=self.pathy[1:]
        dis_y=np.array(y1)-np.array(y0)
        dist=np.sqrt(np.square(dis_x)+np.square(dis_y))
        k= self.T/sum(dist)
        ts=dist*k
        return ts

    def report_planning_result(self):
        if self.plan_server.is_preempt_requested() and self.over_planning_flag ==1:
            rospy.loginfo("Planning preempted!\n")
            self.end_mission()
            self.plan_server.set_preempted()
        else:
            while not self.reached_target:  # when planning is finished, traj tracking will continue to run for a while
                time.sleep(0.01)

            result = PlanResult()
            result.success = True
            self.plan_server.set_succeeded(result)

    def warm_up(self):
        # Send a few setpoints before switching to OFFBOARD mode
        self.state_cmd.position.x = self.drone_state.global_pos[0]
        self.state_cmd.position.y = self.drone_state.global_pos[1]
        self.state_cmd.position.z = self.drone_state.global_pos[2]
        rate = rospy.Rate(100)
        for _ in range(5):  # set 5 points
            if (rospy.is_shutdown()):
                break
            self.local_pos_cmd_pub.publish(self.state_cmd)
            rate.sleep()

    def enter_offboard(self):
        '''
        if not in OFFBOARD mode, switch to OFFBOARD mode
        '''
        if self.flight_state.mode != "OFFBOARD":
            self.warm_up()
            set_offb_req = SetModeRequest()
            set_offb_req.custom_mode = 'OFFBOARD'
            if (self.set_mode_client.call(set_offb_req).mode_sent == True):
                rospy.loginfo("OFFBOARD enabled")

    def start_tracking(self):
        '''
        When triggered, start to publish full state cmd
        '''
        self.enter_offboard()

        self.tracking_cmd_timer = rospy.Timer(rospy.Duration(1/self.cmd_hz), self.tracking_cmd_timer_cb)

    def tracking_cmd_timer_cb(self, event):
        '''
        Publish state cmd, height is fixed to current height
        '''
        self.state_cmd.position.x = self.des_state_array[self.des_state_index][0][0]
        self.state_cmd.position.y = self.des_state_array[self.des_state_index][0][1]
        self.state_cmd.position.z = self.des_pos_z

        self.state_cmd.velocity.x = self.des_state_array[self.des_state_index][1][0]
        self.state_cmd.velocity.y = self.des_state_array[self.des_state_index][1][1]
        self.state_cmd.velocity.z = 0

        self.state_cmd.acceleration_or_force.x = self.des_state_array[self.des_state_index][2][0]
        self.state_cmd.acceleration_or_force.y = self.des_state_array[self.des_state_index][2][1]
        self.state_cmd.acceleration_or_force.z = 0

        self.state_cmd.yaw = np.arctan2(self.des_state_array[self.des_state_index][0][1] - self.des_state_array[self.des_state_index - 1][0][1],
                                        self.des_state_array[self.des_state_index][0][0] - self.des_state_array[self.des_state_index - 1][0][0])

        self.state_cmd.header.stamp = rospy.Time.now()

        self.local_pos_cmd_pub.publish(self.state_cmd)

        if self.des_state_index < len(self.des_state_array) - 1:
            self.des_state_index += 1

    def visualize_obstacle(self):
        obstacle_region = self.visualizer.show_obstacle(self.ox_centre,self.oy_centre,self.des_pos_z,self.obstacle_r)
        self.obstacle_region_pub.publish(obstacle_region)


    def visualize_des_path(self):
        '''
        Visualize the desired path, where high-speed pieces and low-speed pieces are colored differently
        '''
        pos_array = self.planner.get_pos_array()
        pos_array = np.hstack((pos_array, self.des_pos_z * np.ones([len(pos_array), 1])))
        vel_array = np.linalg.norm(self.planner.get_vel_array(), axis=1)
        des_path = self.visualizer.get_path(pos_array, vel_array)
        self.des_path_pub.publish(des_path)


if __name__ == "__main__":

    traj_planner = TrajPlanner()

    rospy.spin()
